/*
 * Copyright 2020-2022 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.naja.auth2server.security;

import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.userdetails.UserDetailsService;
import org.springframework.security.oauth2.client.authentication.OAuth2AuthenticationToken;
import org.springframework.security.oauth2.core.oidc.IdTokenClaimNames;
import org.springframework.security.oauth2.core.oidc.OidcIdToken;
import org.springframework.security.oauth2.core.oidc.endpoint.OidcParameterNames;
import org.springframework.security.oauth2.core.oidc.user.OidcUser;
import org.springframework.security.oauth2.core.user.OAuth2User;
import org.springframework.security.oauth2.jwt.JwtClaimNames;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.token.JwtEncodingContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenCustomizer;

import java.util.*;

/**
 * An {@link OAuth2TokenCustomizer} to map claims from a federated identity to
 * the {@code id_token} produced by this authorization server.
 *
 * @author Steve Riesenberg
 * @since 0.2.3
 */
public final class FederatedIdentityIdTokenCustomizer implements OAuth2TokenCustomizer<JwtEncodingContext> {

    private UserDetailsService userDetailsService;

    public FederatedIdentityIdTokenCustomizer(UserDetailsService userDetailsService) {
        this.userDetailsService = userDetailsService;
    }

    private static final Set<String> ID_TOKEN_CLAIMS = Collections.unmodifiableSet(new HashSet<>(Arrays.asList(
            IdTokenClaimNames.ISS,
            IdTokenClaimNames.SUB,
            IdTokenClaimNames.AUD,
            IdTokenClaimNames.EXP,
            IdTokenClaimNames.IAT,
            IdTokenClaimNames.AUTH_TIME,
            IdTokenClaimNames.NONCE,
            IdTokenClaimNames.ACR,
            IdTokenClaimNames.AMR,
            IdTokenClaimNames.AZP,
            IdTokenClaimNames.AT_HASH,
            IdTokenClaimNames.C_HASH
    )));

    @Override
    public void customize(JwtEncodingContext context) {
        if (OidcParameterNames.ID_TOKEN.equals(context.getTokenType().getValue()) || OAuth2TokenType.ACCESS_TOKEN.getValue().equals(context.getTokenType().getValue())) {
            Authentication principal = context.getPrincipal();
            LoginUser loginUser = null;
            if (principal instanceof UsernamePasswordAuthenticationToken) {
                loginUser = (LoginUser) context.getPrincipal().getPrincipal();
            }
            if (principal instanceof OAuth2AuthenticationToken) {
                OAuth2AuthenticationToken principalOAuth2AuthenticationToken = (OAuth2AuthenticationToken) principal;
                String authorizedClientRegistrationId = principalOAuth2AuthenticationToken.getAuthorizedClientRegistrationId();
                String name = principalOAuth2AuthenticationToken.getName();
                loginUser = (LoginUser) this.userDetailsService.loadUserByUsername(name);
            }
            if (loginUser != null) {
                context.getClaims().claim(JwtClaimNames.JTI, loginUser.getUser().getId());
                context.getClaims().claim("userId", loginUser.getUser().getId());
                context.getClaims().claim("userName", loginUser.getUser().getUserName());
                context.getClaims().claim("mobile", loginUser.getUser().getMobile());
                context.getClaims().claim("email", loginUser.getUser().getEmail());
                context.getClaims().claim("nickName", loginUser.getUser().getNickName());
                context.getClaims().claim("avatarUrl", loginUser.getUser().getAvatarUrl());
                context.getClaims().claim("createTime", loginUser.getUser().getCreateTime());
                context.getClaims().claim("updateTime", loginUser.getUser().getUpdateTime());
                context.getClaims().claim("loginTime", loginUser.getUser().getLoginTime());
                context.getClaims().claim("certificated", loginUser.getUser().getCertificated());
            }
        }
        if (OidcParameterNames.ID_TOKEN.equals(context.getTokenType().getValue())) {
            //idToken存储用户信息,添加一些不重要的属性
        }

    }

    private Map<String, Object> extractClaims(Authentication principal) {
        Map<String, Object> claims;
        if (principal.getPrincipal() instanceof OidcUser) {
            OidcUser oidcUser = (OidcUser) principal.getPrincipal();
            OidcIdToken idToken = oidcUser.getIdToken();
            claims = idToken.getClaims();
        } else if (principal.getPrincipal() instanceof OAuth2User) {
            OAuth2User oauth2User = (OAuth2User) principal.getPrincipal();
            claims = oauth2User.getAttributes();
        } else {
            claims = Collections.emptyMap();
        }

        return new HashMap<>(claims);
    }

}
